module ctrl_unit (
    input   logic   [5:0]       op,
    input   logic   [5:0]       func,


    output  logic               write_reg,
    output  logic               write_mem,
    output  logic               mem2reg,
    output  logic   [2:0]       alu_ctrl,
    output  logic               alu_b_imm,
    output  logic               rb_waddr_rt,
    output  logic               jal,
    output  logic               imm_signedEx,

    input   logic   [4:0]       rs,
    input   logic   [4:0]       rt,
    input   logic   [4:0]       exe_rb_waddr,
    input   logic               exe_mem2reg,
    input   logic               exe_write_reg,
    input   logic   [4:0]       mem_rb_waddr,
    input   logic               mem_mem2reg,
    input   logic               mem_write_reg,

    output  logic   [1:0]       fw_dataA,
    output  logic   [1:0]       fw_dataB,

    input   logic               rs_rt_equ,

    output  logic               write_pc_ir,
    output  logic   [1:0]       pc_source
);


logic write_reg_origin;
logic write_mem_origin;
logic stall;


always_comb begin
    case (op)
        6'b000000: begin    // R_Type
            write_reg_origin    = 1'b1;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b0;
            rb_waddr_rt         = 1'b0;
            jal                 = 1'b0;
            imm_signedEx        = 1'b0;
            case (func)
                6'b100000: alu_ctrl = 3'b100; // +
                6'b100010: alu_ctrl = 3'b101; // -
                6'b100100: alu_ctrl = 3'b000; // &
                6'b100101: alu_ctrl = 3'b001; // |
                6'b100110: alu_ctrl = 3'b010; // ^
                6'b100111: alu_ctrl = 3'b011; // ~|
                6'b101011: alu_ctrl = 3'b110; // < ?1 :0
                6'b000100: alu_ctrl = 3'b111; // <<
                6'b001000: alu_ctrl = 3'b100; // jr
                default:   alu_ctrl = 3'b000;
            endcase
        end
        6'b001000: begin    // addi
            write_reg_origin    = 1'b1;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b1;
            rb_waddr_rt         = 1'b1;
            jal                 = 1'b0;
            imm_signedEx        = 1'b1;
            alu_ctrl            = 3'b100;
        end
        6'b001100: begin    // andi
            write_reg_origin    = 1'b1;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b1;
            rb_waddr_rt         = 1'b1;
            jal                 = 1'b0;
            imm_signedEx        = 1'b0;
            alu_ctrl            = 3'b000;
        end
        6'b001110: begin    // xori
            write_reg_origin    = 1'b1;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b1;
            rb_waddr_rt         = 1'b1;
            jal                 = 1'b0;
            imm_signedEx        = 1'b0;
            alu_ctrl            = 3'b010;
        end
        6'b001011: begin    // sltiu
            write_reg_origin    = 1'b1;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b1;
            rb_waddr_rt         = 1'b1;
            jal                 = 1'b0;
            imm_signedEx        = 1'b0;
            alu_ctrl            = 3'b110;
        end
        6'b100011: begin    // lw
            write_reg_origin    = 1'b1;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b1;
            alu_b_imm           = 1'b1;
            rb_waddr_rt         = 1'b1;
            jal                 = 1'b0;
            imm_signedEx        = 1'b1;
            alu_ctrl            = 3'b100;
        end
        6'b101011: begin    // sw
            write_reg_origin    = 1'b0;
            write_mem_origin    = 1'b1;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b1;
            rb_waddr_rt         = 1'b0;
            jal                 = 1'b0;
            imm_signedEx        = 1'b1;
            alu_ctrl            = 3'b100;
        end
        6'b000100,          // beq
        6'b000101:          // bne
        begin
            write_reg_origin    = 1'b0;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b0;
            rb_waddr_rt         = 1'b0;
            jal                 = 1'b0;
            imm_signedEx        = 1'b0;
            alu_ctrl            = 3'b000;
        end
        6'b000010: begin    // j
            write_reg_origin    = 1'b0;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b0;
            rb_waddr_rt         = 1'b0;
            jal                 = 1'b0;
            imm_signedEx        = 1'b0;
            alu_ctrl            = 3'b000;
        end
        6'b000011: begin    // jal
            write_reg_origin    = 1'b1;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b0;
            rb_waddr_rt         = 1'b0;
            jal                 = 1'b1;
            imm_signedEx        = 1'b0;
            alu_ctrl            = 3'b000;
        end
        default: begin
            write_reg_origin    = 1'b0;
            write_mem_origin    = 1'b0;
            mem2reg             = 1'b0;
            alu_b_imm           = 1'b0;
            rb_waddr_rt         = 1'b0;
            jal                 = 1'b0;
            imm_signedEx        = 1'b0;
            alu_ctrl            = 3'b000;
        end
    endcase
end

always_comb begin
    if((op == 6'b000000) &&(func == 6'b001000)) // jr
        pc_source = 2'd2;
    else
        case (op)
            6'b000100:  // beq
                pc_source = (rs_rt_equ) ?2'd1 :2'd0;
            6'b000101:  // bne
                pc_source = (rs_rt_equ) ?2'd0 :2'd1; 
            6'b000010,  // j
            6'b000011:  // jal
                pc_source = 2'd3;
            default: 
                pc_source = 2'd0;
        endcase
end

always_comb begin
    if(exe_write_reg &&(exe_rb_waddr != 'd0) &&(exe_rb_waddr == rs) &&(!exe_mem2reg))
        fw_dataA = 2'd1;    // exe_alu
    else if(mem_write_reg &&(mem_rb_waddr != 'd0) &&(mem_rb_waddr == rs) &&(!mem_mem2reg))
        fw_dataA = 2'd2;    // mem_alu
    else if(mem_write_reg &&(mem_rb_waddr != 'd0) &&(mem_rb_waddr == rs) &&( mem_mem2reg))
        fw_dataA = 2'd3;    // mem_lw
    else
        fw_dataA = 2'd0;

    if(exe_write_reg &&(exe_rb_waddr != 'd0) &&(exe_rb_waddr == rt) &&(!exe_mem2reg))
        fw_dataB = 2'd1;    // exe_alu
    else if(mem_write_reg &&(mem_rb_waddr != 'd0) &&(mem_rb_waddr == rt) &&(!mem_mem2reg))
        fw_dataB = 2'd2;    // mem_alu
    else if(mem_write_reg &&(mem_rb_waddr != 'd0) &&(mem_rb_waddr == rt) &&( mem_mem2reg))
        fw_dataB = 2'd3;    // mem_lw
    else
        fw_dataB = 2'd0;

    stall = exe_write_reg &&(exe_mem2reg) &&(exe_rb_waddr != 'd0) 
        &&(((rs != 'd0) &&(exe_rb_waddr == rs)) ||((rt != 'd0) &&(exe_rb_waddr == rt)));

    write_pc_ir = ~ stall;
    write_reg   = write_reg_origin & (~ stall);
    write_mem   = write_mem_origin & (~ stall); 
end


endmodule
